// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/exec/aggregate.h>
#include <pollux/expression/function_signature.h>

namespace kumo::pollux {

class AggregateFunc : public exec::Aggregate {
 public:
  explicit AggregateFunc(TypePtr resultType) : Aggregate(resultType) {}

  int32_t accumulatorFixedWidthSize() const override {
    return 0;
  }

  void addRawInput(
      char** /*groups*/,
      const SelectivityVector& /*rows*/,
      const std::vector<VectorPtr>& /*args*/,
      bool /*mayPushdown*/) override {}

  void addIntermediateResults(
      char** /*groups*/,
      const SelectivityVector& /*rows*/,
      const std::vector<VectorPtr>& /*args*/,
      bool /*mayPushdown*/) override {}

  void addSingleGroupRawInput(
      char* /*group*/,
      const SelectivityVector& /*rows*/,
      const std::vector<VectorPtr>& /*args*/,
      bool /*mayPushdown*/) override {}

  void addSingleGroupIntermediateResults(
      char* /*group*/,
      const SelectivityVector& /*rows*/,
      const std::vector<VectorPtr>& /*args*/,
      bool /*mayPushdown*/) override {}

  void extractValues(
      char** /*groups*/,
      int32_t /*numGroups*/,
      VectorPtr* /*result*/) override {}

  void extractAccumulators(
      char** /*groups*/,
      int32_t /*numGroups*/,
      VectorPtr* /*result*/) override {}
  static std::vector<std::shared_ptr<exec::AggregateFunctionSignature>>
  signatures() {
    std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures{
        exec::AggregateFunctionSignatureBuilder()
            .returnType("bigint")
            .intermediateType("array(bigint)")
            .argumentType("bigint")
            .argumentType("double")
            .build(),
        exec::AggregateFunctionSignatureBuilder()
            .typeVariable("T")
            .returnType("T")
            .intermediateType("array(T)")
            .argumentType("T")
            .argumentType("T")
            .build(),
        exec::AggregateFunctionSignatureBuilder()
            .returnType("date")
            .intermediateType("date")
            .build(),
    };
    return signatures;
  }

 protected:
  void initializeNewGroupsInternal(
      char** /*groups*/,
      melon::Range<const vector_size_t*> /*indices*/) override {}
};

inline bool registerAggregateFunc(
    const std::string& name,
    bool overwrite = false) {
  auto signatures = AggregateFunc::signatures();

  return registerAggregateFunction(
             name,
             std::move(signatures),
             [&](core::AggregationNode::Step step,
                 const std::vector<TypePtr>& argTypes,
                 const TypePtr& resultType,
                 const core::QueryConfig& /*config*/)
                 -> std::unique_ptr<exec::Aggregate> {
               if (exec::isPartialOutput(step)) {
                 if (argTypes.empty()) {
                   return std::make_unique<AggregateFunc>(resultType);
                 }
                 return std::make_unique<AggregateFunc>(ARRAY(resultType));
               }
               return std::make_unique<AggregateFunc>(resultType);
             },
             /*registerCompanionFunctions*/ false,
             overwrite)
      .mainFunction;
}

} // namespace kumo::pollux
